package com.cloud.tools.redis.aspect;

import com.cloud.common.core.asserts.ApiAssert;
import com.cloud.common.core.enums.ErrorCodeEnum;
import com.cloud.tools.redis.annotation.RedisLimit;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * 限流
 * 根据不同类创建令牌桶
 * 拦截注解@RedisLimit下的接口，统一限流
 *
 * @author yan
 * @see RedisLimit
 */
@Slf4j
@Aspect
@Component
public class RedisLimitAspect {

    /**
     * 用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter)
     */
    private ConcurrentHashMap<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();

    /**
     * 令牌 限制流量 每秒生成n个令牌 一开始就会初始化n个
     */
    private RateLimiter rateLimiter;

    @Pointcut("@annotation(com.cloud.tools.redis.annotation.RedisLimit)")
    @SuppressWarnings("EmptyMethod")
    public void pointCut() {
    }

    @Around(value = "pointCut()&&@annotation(redisLimit)")
    public Object around(ProceedingJoinPoint point, RedisLimit redisLimit) throws Throwable {
        //获取令牌毫秒数
        long timeout = redisLimit.timeout();

        //获取拦截的方法名
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        //注解所在方法名区分不同的限流策略
        String functionName = method.getDeclaringClass().getName() + "#" + method.getName();
        System.out.println("functionName = " + functionName);
        //不存在就创建
        if (!rateLimiterMap.containsKey(functionName)) {
            rateLimiterMap.put(functionName, RateLimiter.create(redisLimit.limitNum()));
        }
        //获取rateLimiter
        rateLimiter = rateLimiterMap.get(functionName);

        //尝试获取令牌 timout为等待时间 时间过了还未获取到令牌 则抛出异常
        ApiAssert.isTrue(ErrorCodeEnum.INTERNAL_SERVER_ERROR.convert("系统繁忙,请稍候再试"), rateLimiter.tryAcquire(timeout, TimeUnit.MILLISECONDS));

        return point.proceed();

    }

}

